from fastapi import FastAPI, File, UploadFile, HTTPException, Query
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from typing import List, Optional
import logging
from io import BytesIO
from PIL import Image
import base64
from contextlib import asynccontextmanager

from services.ocr_service import OCRService
from services.license_plate_service import LicensePlateService
from models.response_models import (
    Base64ImageRequest, 
    DigitOCRRequest, 
    OCRRecognizeRequest, 
    OCRResponse, 
    HealthResponse,
    LanguagesResponse
)

from config.settings import settings
from utils.file_utils import validate_image_file

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize OCR service
ocr_service = OCRService()
license_plate_service = LicensePlateService()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage the application lifespan"""
    # Startup
    try:
        ocr_service.initialize()
        logger.info("OCR service initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize OCR service: {e}")
        raise e
    
    yield  # Application is running
    
    # Shutdown (if needed)
    logger.info("Shutting down OCR service")

app = FastAPI(
    title="OCR API Service",
    description="OCR image recognition service using PaddleOCR",
    version="1.0.0",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/", response_model=HealthResponse)
async def root():
    """Health check endpoint"""
    return HealthResponse(
        status="healthy",
        message="OCR API Service is running",
        version="1.0.0"
    )

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Detailed health check"""
    try:
        # Test OCR service
        ocr_status = ocr_service.is_healthy()
        return HealthResponse(
            status="healthy" if ocr_status else "unhealthy",
            message="OCR service is ready" if ocr_status else "OCR service is not ready",
            version="1.0.0"
        )
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        return HealthResponse(
            status="unhealthy",
            message=f"Health check failed: {str(e)}",
            version="1.0.0"
        )
@app.post("/api/ocr/recognize", response_model=OCRResponse)
async def recognize_text(
    file: UploadFile = File(..., description="Image file to process"),
    language: str = Query("ch", description="OCR language (ch, en, fr, de, ja, ko, etc.)"),
    confidence_threshold: float = Query(0.5, ge=0.0, le=1.0, description="Minimum confidence threshold for text detection"),
    enhance_image: bool = Query(False, description="Apply image enhancement (grayscale, denoising, contrast adjustment)"),
    correct_skew: bool = Query(False, description="Apply skew correction using Hough transform")
):
    """Advanced OCR recognition with preprocessing options
    
    This endpoint performs OCR on uploaded images with optional preprocessing steps:
    1. Image Enhancement: Grayscale conversion, median filtering, CLAHE contrast adjustment, and sharpening
    2. Skew Correction: Hough transform-based skew detection and correction
    
    Args:
        file: Image file (JPG, PNG, BMP, TIFF, WebP)
        language: OCR language setting
        confidence_threshold: Minimum confidence threshold for text detection (0.0-1.0)
        enhance_image: Apply image enhancement techniques
        correct_skew: Apply skew correction using Hough transform
        
    Returns:
        OCRResponse with text recognition results
    """
    try:
        # Validate file type
        if not file.content_type or not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="File must be an image")
        
        # Validate file size (10MB limit)
        if file.size and file.size > settings.MAX_FILE_SIZE:
            raise HTTPException(
                status_code=400, 
                detail=f"File size too large. Maximum size: {settings.MAX_FILE_SIZE} bytes"
            )
        
        # Validate confidence threshold
        if not 0.0 <= confidence_threshold <= 1.0:
            raise HTTPException(
                status_code=400, 
                detail="Confidence threshold must be between 0.0 and 1.0"
            )
        
        # Read image data
        image_data = await file.read()
        
        # Validate image data
        if len(image_data) == 0:
            raise HTTPException(status_code=400, detail="Empty image file")
        
        logger.info(f"Processing OCR request: file={file.filename}, language={language}, confidence_threshold={confidence_threshold}")
        
        # Perform OCR with preprocessing options
        result = ocr_service.recognize_image(
            image_data=image_data,
            language=language,
            confidence_threshold=confidence_threshold,
            enhance=enhance_image,
            correct_skew=correct_skew
        )
        
        return OCRResponse(
            success=True,
            message=f"OCR recognition completed successfully. Found {len(result.results)} text regions.",
            data=result
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Enhanced OCR processing failed: {e}")
        raise HTTPException(status_code=500, detail=f"OCR processing failed: {str(e)}")

@app.post("/ocr/upload", response_model=OCRResponse)
async def ocr_from_upload(file: UploadFile = File(...)):
    """OCR recognition from uploaded image file (legacy endpoint)"""
    try:
        # Validate file type
        if not file.content_type or not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="File must be an image")
        
        # Read image data
        image_data = await file.read()
        
        # Perform OCR with default settings
        result = ocr_service.recognize_image(image_data)
        
        return OCRResponse(
            success=True,
            message="OCR recognition completed successfully",
            data=result
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"OCR upload processing failed: {e}")
        raise HTTPException(status_code=500, detail=f"OCR processing failed: {str(e)}")

@app.post("/ocr/base64", response_model=OCRResponse)
async def ocr_from_base64(request: Base64ImageRequest):
    """OCR recognition from base64 encoded image"""
    try:
        # Decode base64 image
        try:
            image_base64 = request.image_base64
            # Remove data URL prefix if present
            if image_base64.startswith('data:image'):
                image_base64 = image_base64.split(',')[1]
            
            image_data = base64.b64decode(image_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid base64 image data: {str(e)}")
        
        # Perform OCR
        result = ocr_service.recognize_image(image_data)
        
        return OCRResponse(
            success=True,
            message="OCR recognition completed successfully",
            data=result
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"OCR base64 processing failed: {e}")
        raise HTTPException(status_code=500, detail=f"OCR processing failed: {str(e)}")

@app.get("/ocr/languages", response_model=LanguagesResponse)
async def get_supported_languages():
    """Get list of supported OCR languages and model information
    
    Returns:
        LanguagesResponse with supported languages and model information
    """
    try:
        # Get model information
        model_info = {
            "detection_model": "ch_PP-OCRv4_det_infer",
            "recognition_model": "ch_PP-OCRv4_rec_infer",
            "classification_model": "ch_ppocr_mobile_v2.0_cls_infer",
            "high_accuracy_available": True,
            "high_accuracy_models": {
                "detection": "ch_PP-OCRv4_det_server_infer",
                "recognition": "ch_PP-OCRv4_rec_server_infer"
            }
        }
        
        if settings.OCR_USE_HIGH_ACCURACY:
            model_info["detection_model"] = "ch_PP-OCRv4_det_server_infer"
            model_info["recognition_model"] = "ch_PP-OCRv4_rec_server_infer"
        
        return LanguagesResponse(
            success=True,
            message="Supported languages and model information retrieved successfully",
            data={
                "languages": ["ch", "en", "fr", "german", "korean", "japan", "chinese_cht"],
                "default_language": "ch",
                "model_info": model_info
            }
        )
    except Exception as e:
        logger.error(f"Failed to retrieve supported languages: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to retrieve supported languages: {str(e)}")

@app.post("/api/ocr/digits", response_model=OCRResponse)
async def recognize_digits(
    file: UploadFile = File(..., description="Image file containing digits to process"),
    confidence_threshold: float = Query(0.8, ge=0.0, le=1.0, description="Minimum confidence threshold for digit detection"),
    enhance_image: bool = Query(False, description="Apply image enhancement (grayscale, denoising, contrast adjustment)"),
    correct_skew: bool = Query(False, description="Apply skew correction using Hough transform")
):
    """OCR recognition optimized for digits only with preprocessing options
    
    This endpoint performs OCR optimized for digit recognition with optional preprocessing:
    1. Image Enhancement: Grayscale conversion, median filtering, CLAHE contrast adjustment, and sharpening
    2. Skew Correction: Hough transform-based skew detection and correction
    
    Args:
        file: Image file (JPG, PNG, BMP, TIFF, WebP)
        confidence_threshold: Minimum confidence threshold for digit detection (0.0-1.0)
        enhance_image: Apply image enhancement techniques
        correct_skew: Apply skew correction using Hough transform
        
    Returns:
        OCRResponse with digit recognition results
    """
    try:
        # Validate file type
        if not file.content_type or not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="File must be an image")
        
        # Validate file size (10MB limit)
        if file.size and file.size > settings.MAX_FILE_SIZE:
            raise HTTPException(
                status_code=400, 
                detail=f"File size too large. Maximum size: {settings.MAX_FILE_SIZE} bytes"
            )
        
        # Validate confidence threshold
        if not 0.0 <= confidence_threshold <= 1.0:
            raise HTTPException(
                status_code=400, 
                detail="Confidence threshold must be between 0.0 and 1.0"
            )
        
        # Read image data
        image_data = await file.read()
        
        # Validate image data
        if len(image_data) == 0:
            raise HTTPException(status_code=400, detail="Empty image file")
        
        logger.info(f"Processing digit-only OCR request: file={file.filename}, confidence_threshold={confidence_threshold}")
        
        # Perform digit-only OCR with preprocessing options
        result = ocr_service.recognize_digits_only(
            image_data=image_data,
            confidence_threshold=confidence_threshold,
            enhance=enhance_image,
            correct_skew=correct_skew
        )
        
        return OCRResponse(
            success=True,
            message=f"Digit-only OCR recognition completed successfully. Found {len(result.results)} digit regions.",
            data=result
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Digit-only OCR processing failed: {e}")
        raise HTTPException(status_code=500, detail=f"OCR processing failed: {str(e)}")

@app.post("/api/ocr/digits_without_license_plate", response_model=OCRResponse)
async def recognize_digits_without_license_plate(
    file: UploadFile = File(..., description="Image file containing digits to process (license plate regions will be removed)"),
    confidence_threshold: float = 0.8,
    remove_duplicates: bool = True,
    enhance_image: bool = Query(False, description="Apply image enhancement (grayscale, denoising, contrast adjustment)"),
    correct_skew: bool = Query(False, description="Apply skew correction using Hough transform")
):
    """OCR recognition optimized for digits only, with license plate regions removed or blurred
    
    This endpoint detects license plate regions in the image, blurs them, and then recognizes
    digits outside of those regions. Optionally removes duplicate digits from results.
    
    Args:
        file: Image file (JPG, PNG, BMP, TIFF, WebP)
        confidence_threshold: Minimum confidence threshold for digit detection (0.0-1.0)
        remove_duplicates: Whether to remove duplicate digits from results (e.g., 123123123 becomes 123)
        enhance_image: Apply image enhancement techniques
        correct_skew: Apply skew correction using Hough transform
    
    Returns:
        OCRResponse with digit recognition results excluding license plate regions
    """
    try:
        # Validate file type
        if not file.content_type or not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="File must be an image")
        
        # Validate file size (10MB limit)
        if file.size and file.size > settings.MAX_FILE_SIZE:
            raise HTTPException(
                status_code=400, 
                detail=f"File size too large. Maximum size: {settings.MAX_FILE_SIZE} bytes"
            )
        
        # Validate confidence threshold
        if not 0.0 <= confidence_threshold <= 1.0:
            raise HTTPException(
                status_code=400, 
                detail="Confidence threshold must be between 0.0 and 1.0"
            )
        
        # Read image data
        image_data = await file.read()
        
        # Validate image data
        if len(image_data) == 0:
            raise HTTPException(status_code=400, detail="Empty image file")
        
        logger.info(f"Processing digits OCR without license plate: file={file.filename}, confidence_threshold={confidence_threshold}")
        
        # Perform OCR with license plate removal and preprocessing options
        result = license_plate_service.process_image_without_license_plate(
            image_data=image_data,
            confidence_threshold=confidence_threshold,
            remove_duplicates=remove_duplicates,
            enhance=enhance_image,
            correct_skew=correct_skew
        )
        
        return OCRResponse(
            success=True,
            message=f"Digit OCR without license plate completed successfully. Found {len(result.results)} digit regions.",
            data=result
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Digit OCR without license plate processing failed: {e}")
        raise HTTPException(status_code=500, detail=f"OCR processing failed: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host=settings.HOST,
        port=settings.PORT,
        reload=settings.DEBUG,
        access_log=True
    )